﻿/*
** Xin YUAN, 2023, BSD (2)
*/

////////////////////////////////////////////////////////////////////////////////

//_w_child_process
/*! \brief A class for sub-process.

A class for sub-process.
*/
class _w_child_process
{
public:
	/*! \brief Constructor.

	Constructor.
	*/
	_w_child_process() noexcept
	{
		clear_info();
	}
	_w_child_process(const _w_child_process&) = delete;
	_w_child_process(_w_child_process&& src) noexcept
	{
		m_siStartupInfo = src.m_siStartupInfo;
		m_piProcessInfo = src.m_piProcessInfo;
		src.clear_info();
	}
	/*! \brief Destructor.

	Destructor.
	*/
	~_w_child_process() noexcept
	{
		CloseProcessHandles();
	}

	_w_child_process& operator=(const _w_child_process&) = delete;
	_w_child_process& operator=(_w_child_process&& src) noexcept
	{
		if ( this != &src ) {
			assert( m_piProcessInfo.hProcess == NULL && m_piProcessInfo.hThread == NULL );
			//CloseProcessHandles();
			m_siStartupInfo = src.m_siStartupInfo;
			m_piProcessInfo = src.m_piProcessInfo;
			src.clear_info();
		}
		return *this;
	}

	/*! \brief Close all handles.

	Close all handles.
	*/
	void CloseProcessHandles() noexcept
	{
		//Close Handles
		BOOL bRes = TRUE;
		(void)bRes;
		if ( m_piProcessInfo.hProcess != NULL ) {
			bRes = ::CloseHandle(m_piProcessInfo.hProcess);
			assert( bRes );
			m_piProcessInfo.hProcess = NULL;
		}
		if ( m_piProcessInfo.hThread != NULL ) {
			bRes = ::CloseHandle(m_piProcessInfo.hThread);
			assert( bRes );
			m_piProcessInfo.hThread = NULL;
		}
	}

	/*! \brief Get startup information.

	Get startup information.
	\return A reference to STARTUPINFOW.
	*/
	const STARTUPINFOW& GetStartupInfo() const noexcept
	{
		return m_siStartupInfo;
	}
	STARTUPINFOW& GetStartupInfo() noexcept
	{
		return m_siStartupInfo;
	}
	/*! \brief Get process information.

	Get process information.
	\return A reference to PROCESS_INFORMATION.
	*/
	const PROCESS_INFORMATION& GetProcessInfo() const noexcept
	{
		return m_piProcessInfo;
	}
	PROCESS_INFORMATION& GetProcessInfo() noexcept
	{
		return m_piProcessInfo;
	}

	/*! \brief Get the handle of child process.

	Get the handle of child process.
	\return A handle of child process.
	*/
	HANDLE GetChildProcessHandle() const noexcept
	{
		return m_piProcessInfo.hProcess;
	}
	/*! \brief Get the handle of child thread.

	Get the handle of child thread.
	\return A handle of child thread.
	*/
	HANDLE GetChildThreadHandle() const noexcept
	{
		return m_piProcessInfo.hThread;
	}

	void SetStartupFlags(BOOL bShowWindow) noexcept
	{
		m_siStartupInfo.dwFlags |= STARTF_USESHOWWINDOW;
		m_siStartupInfo.wShowWindow = bShowWindow ? SW_SHOWNORMAL : SW_HIDE;
	}

	/*! \brief Create a child process.

	Create a child process.
	\param szCmdLine [in] Specify the command line string.
	\param bInherit [in] Specify the handles should be inherited.
	\param szWorkPath [in] Specify the work path of child process.
	\param dwFlags [in] Specify the flags. See CreateProcessW().
	                    Some useful falgs are CREATE_NEW_CONSOLE, CREATE_SUSPENDED, CREATE_NO_WINDOW.
	\return TRUE for succeeded, FALSE for otherwise.
	*/
	BOOL Create(LPCWSTR szCmdLine, BOOL bInherit,
			LPCWSTR szWorkPath = NULL, DWORD dwFlags = 0) noexcept
	{
		assert( m_piProcessInfo.hProcess == NULL && m_piProcessInfo.hThread == NULL );

		//command line
		WCHAR* szBuf = ::_wcsdup(szCmdLine);
		if ( szBuf == NULL )
			return FALSE;

		dwFlags |= CREATE_UNICODE_ENVIRONMENT;
		// Create the child process.
		BOOL bRet = ::CreateProcessW(NULL,
					szBuf,             // command line
					NULL,              // process security attributes
					NULL,              // primary thread security attributes
					bInherit,          // handles are inherited
					dwFlags,           // creation flags
					NULL,              // use parent's environment
					szWorkPath,        // if NULL, use parent's current directory
					&m_siStartupInfo,  // STARTUPINFOW pointer
					&m_piProcessInfo   // receives PROCESS_INFORMATION
					);
		if ( !bRet ) {
			//FALSE, ::GetLastError()
			::free(szBuf);
			return FALSE;
		} //end if

		::free(szBuf);
		return bRet;
	}

	// for GUI child-process
	// return: 0, WAIT_TIMEOUT, WAIT_FAILED
	DWORD WaitForIdle(DWORD dwMilliseconds = INFINITE) noexcept
	{
		assert( m_piProcessInfo.hProcess != NULL );
		return ::WaitForInputIdle(m_piProcessInfo.hProcess, dwMilliseconds);
	}

	BOOL Resume() noexcept
	{
		assert( m_piProcessInfo.hThread != NULL );
		return ::ResumeThread(m_piProcessInfo.hThread) != (DWORD)-1;  //-1, ::GetLastError()
	}

	/*! \brief Get the exit code of the process.

	Get the exit code of the process.
	\return The exit code of the process or STILL_ACTIVE.
	*/
	DWORD GetExitCode() const noexcept
	{
		assert( m_piProcessInfo.hProcess != NULL );
		DWORD dwCode = (DWORD)-1;
		BOOL bRet = ::GetExitCodeProcess(m_piProcessInfo.hProcess, &dwCode);
		(void)bRet;  //FALSE, ::GetLastError()
		return dwCode;
	}

	/*! \brief Get process state.

	Get process state.
	\return A value for process state, WAIT_OBJECT_0, WAIT_TIMEOUT, WAIT_FAILED.
	*/
	DWORD GetState() const noexcept
	{
		assert( m_piProcessInfo.hProcess != NULL );
		return ::WaitForSingleObject(m_piProcessInfo.hProcess, 0);
	}

	/*! \brief Terminate child process in a specified time.

	Terminate child process in a specified time.
	\param dwWaitTime [in] Specify the waiting time.
	*/
	void WaitForEnd(DWORD dwWaitTime = INFINITE) noexcept
	{
		assert( m_piProcessInfo.hProcess != NULL );
		DWORD dwState = ::WaitForSingleObject(m_piProcessInfo.hProcess, dwWaitTime);
		if ( dwState == WAIT_FAILED || dwState == WAIT_OBJECT_0 )
			return ;
		BOOL bRet = ::TerminateProcess(m_piProcessInfo.hProcess, 0);
		(void)bRet;  //FALSE, ::GetLastError()
		dwState = ::WaitForSingleObject(m_piProcessInfo.hProcess, dwWaitTime);
		if ( dwState == WAIT_TIMEOUT )
			::WaitForSingleObject(m_piProcessInfo.hProcess, INFINITE);  //no check
	}

	void Terminate() noexcept
	{
		assert( m_piProcessInfo.hProcess != NULL );
		::PostThreadMessageW(m_piProcessInfo.dwThreadId, WM_QUIT, 0, 0);  //no check
	}

protected:
	/*! \brief Clear information.

	Clear information.
	*/
	void clear_info() noexcept
	{
		::ZeroMemory(&m_siStartupInfo, sizeof(STARTUPINFOW));
		m_siStartupInfo.cb = sizeof(STARTUPINFOW);
		::ZeroMemory(&m_piProcessInfo, sizeof(PROCESS_INFORMATION));
	}

protected:
	STARTUPINFOW        m_siStartupInfo;  //!< startup information.
	PROCESS_INFORMATION m_piProcessInfo;  //!< process information.
};

//_w_rio_process
/*! \brief A class for redirecting IO process.

A class for redirecting IO process.
*/
class _w_rio_process : public _w_child_process
{
private:
	typedef _w_child_process  baseClass;

public:
	/*! \brief Constructor.

	Constructor.
	*/
	_w_rio_process() noexcept : m_hChildStdoutRd(NULL), m_hChildStdinWr(NULL)
	{
	}
	_w_rio_process(const _w_rio_process&) = delete;
	_w_rio_process(_w_rio_process&& src) noexcept : baseClass(rv_forward(static_cast<baseClass&>(src))),
							m_hChildStdoutRd(src.m_hChildStdoutRd),
							m_hChildStdinWr(src.m_hChildStdinWr)
	{
		src.m_hChildStdoutRd = NULL;
		src.m_hChildStdinWr = NULL;
	}
	/*! \brief Destructor.

	Destructor.
	*/
	~_w_rio_process() noexcept
	{
		ClosePipeStdout();
		ClosePipeStdin();
	}

	_w_rio_process& operator=(const _w_rio_process&) = delete;
	_w_rio_process& operator=(_w_rio_process&& src) noexcept
	{
		baseClass::operator=(rv_forward(static_cast<baseClass&>(src)));
		if ( this != &src ) {
			ClosePipeStdout();
			ClosePipeStdin();
			m_hChildStdoutRd = src.m_hChildStdoutRd;
			m_hChildStdinWr = src.m_hChildStdinWr;
			src.m_hChildStdoutRd = NULL;
			src.m_hChildStdinWr = NULL;
		}
		return *this;
	}

	/*! \brief Close stdout pipe of child process.

	Close stdout pipe of child process.
	*/
	void ClosePipeStdout() noexcept
	{
		BOOL bRes = TRUE;
		(void)bRes;
		if ( m_hChildStdoutRd != NULL ) {
			bRes = ::CloseHandle(m_hChildStdoutRd);
			assert( bRes );
			m_hChildStdoutRd = NULL;
		}
	}
	/*! \brief Close stdin pipe of child process.

	Close stdin pipe of child process.
	*/
	void ClosePipeStdin() noexcept
	{
		BOOL bRes = TRUE;
		(void)bRes;
		if ( m_hChildStdinWr != NULL ) {
			bRes = ::CloseHandle(m_hChildStdinWr);
			assert( bRes );
			m_hChildStdinWr = NULL;
		}
	}

	/*! \brief Create a redirected IO child process.

	Create a redirected IO child process.
	\param szCmdLine [in] Specify the command line string.
	\param szWorkPath [in] Specify the work path of child process.
	\param dwFlags [in] Specify the flags. See CreateProcessW().
	\param dwSize [in] Specify the pipe size. See CreatePipe().
	\return TRUE for succeeded, FALSE for otherwise.
	*/
	BOOL Create(LPCWSTR szCmdLine, LPCWSTR szWorkPath = NULL,
			DWORD dwFlags = 0, DWORD dwSize = 0) noexcept
	{
		assert( m_hChildStdoutRd == NULL && m_hChildStdinWr == NULL );

		HANDLE hChildStdoutRd = NULL;
		HANDLE hChildStdoutWr = NULL;
		HANDLE hChildStdinRd = NULL;
		HANDLE hChildStdinWr = NULL;

		BOOL bRes = TRUE;
		(void)bRes;

		// Set the bInheritHandle flag so pipe handles are inherited.
		SECURITY_ATTRIBUTES saAttr;
		saAttr.nLength = sizeof(SECURITY_ATTRIBUTES);
		saAttr.bInheritHandle = TRUE;
		saAttr.lpSecurityDescriptor = NULL;

		// Create a pipe for the child process's STDOUT.
		if ( !::CreatePipe(&hChildStdoutRd, &hChildStdoutWr, &saAttr, dwSize) )
			return FALSE;  //::GetLastError()

		// Ensure the read handle to the pipe for STDOUT is not inherited.
		if ( !::SetHandleInformation(hChildStdoutRd, HANDLE_FLAG_INHERIT, 0) ) {
			//FALSE, ::GetLastError()
			bRes = ::CloseHandle(hChildStdoutRd);
			assert( bRes );
			bRes = ::CloseHandle(hChildStdoutWr);
			assert( bRes );
			return FALSE;
		}

		// Create a pipe for the child process's STDIN.
		if ( !::CreatePipe(&hChildStdinRd, &hChildStdinWr, &saAttr, dwSize) ) {
			//FALSE, ::GetLastError()
			bRes = ::CloseHandle(hChildStdoutRd);
			assert( bRes );
			bRes = ::CloseHandle(hChildStdoutWr);
			assert( bRes );
			return FALSE;
		}

		// Ensure the write handle to the pipe for STDIN is not inherited.
 		if ( !::SetHandleInformation(hChildStdinWr, HANDLE_FLAG_INHERIT, 0) ) {
			//FALSE, ::GetLastError()
			bRes = ::CloseHandle(hChildStdoutRd);
			assert( bRes );
			bRes = ::CloseHandle(hChildStdoutWr);
			assert( bRes );
			bRes = ::CloseHandle(hChildStdinRd);
			assert( bRes );
			bRes = ::CloseHandle(hChildStdinWr);
			assert( bRes );
			return FALSE;
		}

		// Create the child process.
		m_siStartupInfo.dwFlags |= STARTF_USESTDHANDLES;
		m_siStartupInfo.hStdInput  = hChildStdinRd;
		m_siStartupInfo.hStdOutput = hChildStdoutWr;
		m_siStartupInfo.hStdError  = hChildStdoutWr;
		if ( !baseClass::Create(szCmdLine, TRUE, szWorkPath, dwFlags) ) {
			bRes = ::CloseHandle(hChildStdoutRd);
			assert( bRes );
			bRes = ::CloseHandle(hChildStdoutWr);
			assert( bRes );
			bRes = ::CloseHandle(hChildStdinRd);
			assert( bRes );
			bRes = ::CloseHandle(hChildStdinWr);
			assert( bRes );
			return FALSE;
		}

		// Close handles to the stdin and stdout pipes no longer needed by the child process.
		bRes = ::CloseHandle(hChildStdoutWr);
		assert( bRes );
		bRes = ::CloseHandle(hChildStdinRd);
		assert( bRes );

		m_hChildStdoutRd = hChildStdoutRd;
		m_hChildStdinWr  = hChildStdinWr;
		return TRUE;
	}

	/*! \brief Write data to pipe.

	Write data to pipe.
	\param pData [in] A pointer to buffer.
	\param dwSize [in] The bytes of buffer.
	\param lpWritten [out] Receive the written bytes. It can be NULL.
	\return TRUE for succeeded, FALSE for otherwise.
	\note This method may be blocked.
	*/
	BOOL WriteToPipe(LPCVOID pData, DWORD dwSize, LPDWORD lpWritten) noexcept
	{
		assert( m_hChildStdinWr != NULL );
		return ::WriteFile(m_hChildStdinWr, pData, dwSize,
					lpWritten, NULL);  //FALSE, ::GetLastError()
	}

	/*! \brief Read data from pipe.

	Read data from pipe.
	\param pData [in] A pointer to buffer.
	\param dwSize [in] The bytes of buffer.
	\param lpRead [out] Receive the read bytes. It can be NULL.
	\return TRUE for succeeded, FALSE for otherwise.
	\note This method may be blocked.
	      If *lpRead == 0, the write end of pipe is closed.
	*/
	BOOL ReadFromPipe(LPVOID pData, DWORD dwSize, LPDWORD lpRead) noexcept
	{
		assert( m_hChildStdoutRd != NULL );
		BOOL bRet = ::ReadFile(m_hChildStdoutRd, pData, dwSize,
					lpRead, NULL);
		if ( !bRet ) {
			//FALSE, ::GetLastError()
			if ( ::GetLastError() == ERROR_BROKEN_PIPE ) { //write end closed
				if ( lpRead != NULL )
					*lpRead = 0;
				bRet = TRUE;
			}
		}
		return bRet;
	}

protected:
	HANDLE m_hChildStdoutRd;  //!< Read end for non-inherited child stdout.
	HANDLE m_hChildStdinWr;   //!< Write end for non-inherited child stdin.
};

////////////////////////////////////////////////////////////////////////////////
